import rarfile
import os
import sys
import time
import datetime

class RARViewer:
    def __init__(self, rar_path):
        self.rar_path = rar_path
        self.current_path = ""  # 记录当前浏览路径
        print(f"正在打开RAR文件: {rar_path}")
        print("读取文件目录结构中...")
        
        # 记录开始时间
        start_time = time.time()
        try:
            self.rar = rarfile.RarFile(rar_path)
            end_time = time.time()
            print(f"成功打开RAR文件，耗时 {end_time - start_time:.2f} 秒")
        except Exception as e:
            print(f"无法打开RAR文件: {e}")
            sys.exit(1)
    
    def format_time(self, time_obj):
        """格式化时间对象"""
        if time_obj is None:
            return "N/A"
        
        try:
            # 如果是time.struct_time对象（元组）
            if isinstance(time_obj, tuple):
                dt = datetime.datetime(*time_obj[:6])
                return dt.strftime("%Y-%m-%d %H:%M:%S")
            # 如果有strftime方法（datetime对象）
            elif hasattr(time_obj, 'strftime'):
                return time_obj.strftime("%Y-%m-%d %H:%M:%S")
            else:
                return "N/A"
        except Exception:
            return "N/A"
    
    def list_files(self, start_index=0, count=50):
        """列出文件，支持分页
        Args:
            start_index: 开始索引（从0开始）
            count: 显示文件数量
        """
        print(f"\nRAR文件 '{self.rar_path}' 当前路径 '{self.current_path if self.current_path else '/'}' 的内容:")
        print("-" * 100)
        print(f"{'序号':<5} {'名称':<55} {'大小':<15} {'修改时间':<20}")
        print("-" * 100)
        
        try:
            infos = self.rar.infolist()
            
            # 过滤出当前路径下的项目
            filtered_infos = []
            dir_prefix = self.current_path + "/" if self.current_path else ""
            
            for info in infos:
                # 如果在根目录，只显示顶层项目
                if not self.current_path:
                    if info.filename.count('/') == 0 or (info.filename.count('/') == 1 and info.isdir()):
                        filtered_infos.append(info)
                # 如果在子目录，只显示该目录下的直接子项目
                else:
                    if info.filename.startswith(dir_prefix):
                        relative_path = info.filename[len(dir_prefix):]
                        # 只包含直接子级（不包含子目录的子级）
                        if relative_path and (relative_path.count('/') == 0 or (relative_path.count('/') == 1 and info.isdir())):
                            filtered_infos.append(info)
            
            total_files = len(filtered_infos)
            
            # 处理边界情况
            if start_index >= total_files:
                print(f"起始索引 {start_index} 超出范围。总共只有 {total_files} 个文件。")
                return
            
            if start_index < 0:
                start_index = 0
                
            end_index = min(start_index + count, total_files)
            
            print(f"显示 {start_index}-{end_index-1} 个项目，总共 {total_files} 个文件/目录")
            
            for i in range(start_index, end_index):
                info = filtered_infos[i]
                size = info.file_size
                modified = self.format_time(info.date_time) if hasattr(info, 'date_time') else "N/A"
                
                # 获取相对路径（去掉当前路径前缀）
                if self.current_path:
                    name = info.filename[len(self.current_path) + 1:]
                else:
                    name = info.filename
                
                # 计算缩进级别以显示层次结构
                level = name.count('/') if not info.isdir() else name.rstrip('/').count('/')
                indent = "  " * level
                
                # 如果是目录，在名称后添加斜杠
                if info.isdir():
                    name = name.rstrip('/') + '/'
                
                # 只显示文件名部分（不包括路径）
                basename = os.path.basename(name.rstrip('/')) + ('/' if info.isdir() else '')
                
                print(f"{i:<5} {indent}{basename:<{55-2*level}} {size:<15} {modified:<20}")
                    
            # 显示导航提示
            if end_index < total_files:
                next_start = end_index
                remaining = total_files - end_index
                print(f"\n还有 {remaining} 个文件未显示。")
                print(f"输入 'list {next_start} {count}' 查看下一批文件。")
            else:
                print(f"\n已显示所有 {total_files} 个文件。")
                
        except Exception as e:
            print(f"读取文件列表时出错: {e}")
    
    def list_tree(self, index=None, display_mode="head", count=10):
        """显示指定目录下的文件树形结构
        Args:
            index: 目录的索引号，None表示当前目录
            display_mode: 显示模式 - "head"(前几个), "all"(全部), "tail"(后几个), "size"(显示文件数量)
            count: 在head或tail模式下显示的项目数量
        """
        try:
            infos = self.rar.infolist()
            
            # 如果提供了索引，找到对应的目录
            directory = self.current_path
            if index is not None:
                # 获取当前目录下的项目列表，用于查找索引对应的目录
                current_dir_items = []
                dir_prefix = directory + "/" if directory else ""
                
                for info in infos:
                    # 如果在根目录，只显示顶层项目
                    if not directory:
                        if info.filename.count('/') == 0 or (info.filename.count('/') == 1 and info.isdir()):
                            current_dir_items.append(info)
                    # 如果在子目录，只显示该目录下的直接子项目
                    else:
                        if info.filename.startswith(dir_prefix):
                            relative_path = info.filename[len(dir_prefix):]
                            # 只包含直接子级（不包含子目录的子级）
                            if relative_path and (relative_path.count('/') == 0 or (relative_path.count('/') == 1 and info.isdir())):
                                current_dir_items.append(info)
                
                if index < 0 or index >= len(current_dir_items):
                    print(f"错误: 索引 {index} 超出范围，有效范围为 0-{len(current_dir_items)-1}")
                    return
                
                info = current_dir_items[index]
                if not info.isdir():
                    print(f"错误: 索引 {index} 对应的项目 '{info.filename}' 不是目录")
                    return
                
                # 更新要查看的目录为索引指定的子目录
                directory = info.filename.rstrip('/')
            
            print(f"\nRAR文件 '{self.rar_path}' 当前路径 '{directory if directory else '/'}' 的内容:")
            print("-" * 100)
            
            # 过滤出指定目录下的项目
            filtered_infos = []
            dir_prefix = directory + "/" if directory else ""
            
            for info in infos:
                # 如果指定了目录，只显示该目录下的内容
                if directory:
                    if info.filename.startswith(dir_prefix):
                        relative_path = info.filename[len(dir_prefix):]
                        # 只包含直接子级（不包含子目录的子级）
                        if relative_path and (relative_path.count('/') == 0 or (relative_path.count('/') == 1 and info.isdir())):
                            filtered_infos.append(info)
                else:
                    # 根目录，只显示顶层项目
                    if info.filename.count('/') == 0 or (info.filename.count('/') == 1 and info.isdir()):
                        filtered_infos.append(info)
            
            # 根据显示模式处理
            if display_mode.lower() == "size":
                # 显示文件和目录数量
                total_files = sum(1 for info in filtered_infos if not info.isdir())
                total_dirs = sum(1 for info in filtered_infos if info.isdir())
                print(f"目录统计信息:")
                print(f"  目录数量: {total_dirs}")
                print(f"  文件数量: {total_files}")
                print(f"  总计数量: {len(filtered_infos)}")
                return
            elif display_mode.lower() == "head":
                display_infos = filtered_infos[:count]
                print(f"显示前 {min(count, len(filtered_infos))} 个项目，共 {len(filtered_infos)} 个项目")
            elif display_mode.lower() == "tail":
                display_infos = filtered_infos[-count:]
                print(f"显示后 {min(count, len(filtered_infos))} 个项目，共 {len(filtered_infos)} 个项目")
            else:  # all
                display_infos = filtered_infos
                print(f"显示全部 {len(filtered_infos)} 个项目")
            
            if not filtered_infos:
                print(f"目录 '{directory if directory else '/'}' 为空")
                return
            
            # 构建并显示树形结构
            tree = {}
            # 创建一个映射，将文件名映射到全局索引
            name_to_global_index = {info.filename: i for i, info in enumerate(infos)}
            
            for info in display_infos:
                if directory:
                    relative_name = info.filename[len(dir_prefix):]
                else:
                    relative_name = info.filename
                
                path_parts = relative_name.rstrip('/').split('/')
                current = tree
                for i, part in enumerate(path_parts):
                    is_last = (i == len(path_parts) - 1)
                    if part not in current:
                        current[part] = {
                            "__is_dir__": not is_last or info.isdir(), 
                            "__info__": info if is_last else None, 
                            "__children__": {},
                            "__global_index__": name_to_global_index.get(info.filename, -1) if is_last else -1
                        }
                    current = current[part]["__children__"]
            
            # 显示树形结构
            def print_tree(node, prefix=""):
                keys = list(node.keys())
                for i, key in enumerate(keys):
                    is_last = (i == len(keys) - 1)
                    item = node[key]
                    is_dir = item["__is_dir__"]
                    
                    # 构造连接符
                    connector = "└── " if is_last else "├── "
                    icon = "📁 " if is_dir else "📄 "
                    
                    # 显示文件信息
                    info = item["__info__"]
                    if info:
                        size_str = str(info.file_size) if not is_dir else "-"
                        modified = self.format_time(info.date_time) if hasattr(info, 'date_time') else "N/A"
                        global_index = item["__global_index__"]
                        index_str = f"[{global_index}]" if global_index >= 0 else ""
                        name_with_icon = f"{icon}{key}{('/' if is_dir else '')} {index_str}"
                        print(f"{prefix}{connector}{name_with_icon:<40} {size_str:<12} {modified}")
                    else:
                        name_with_icon = f"{icon}{key}/"
                        print(f"{prefix}{connector}{name_with_icon}")
                    
                    # 为子节点准备前缀
                    extension = "    " if is_last else "│   "
                    
                    # 递归显示子节点
                    if item["__children__"]:
                        print_tree(item["__children__"], prefix + extension)
            
            print_tree(tree)
                
        except Exception as e:
            print(f"读取文件列表时出错: {e}")

    def cd(self, index=None):
        """更改当前目录
        Args:
            index: 要进入的目录索引，None表示返回上一级目录，-1表示返回根目录
        """
        try:
            infos = self.rar.infolist()
            
            if index is None:
                # 返回上一级目录
                if self.current_path:
                    parent_path = "/".join(self.current_path.split("/")[:-1])
                    self.current_path = parent_path
                    print(f"已返回上级目录，当前路径: '{self.current_path if self.current_path else '/'}'")
                else:
                    print("当前已在根目录")
                return
            elif index == -1:
                # 返回根目录
                self.current_path = ""
                print("已返回根目录")
                return
            
            # 进入指定索引的目录
            current_dir_items = []
            dir_prefix = self.current_path + "/" if self.current_path else ""
            
            for info in infos:
                # 如果在根目录，只显示顶层项目
                if not self.current_path:
                    if info.filename.count('/') == 0 or (info.filename.count('/') == 1 and info.isdir()):
                        current_dir_items.append(info)
                # 如果在子目录，只显示该目录下的直接子项目
                else:
                    if info.filename.startswith(dir_prefix):
                        relative_path = info.filename[len(dir_prefix):]
                        # 只包含直接子级（不包含子目录的子级）
                        if relative_path and (relative_path.count('/') == 0 or (relative_path.count('/') == 1 and info.isdir())):
                            current_dir_items.append(info)
            
            if index < 0 or index >= len(current_dir_items):
                print(f"错误: 索引 {index} 超出范围，有效范围为 0-{len(current_dir_items)-1}")
                return
            
            info = current_dir_items[index]
            if not info.isdir():
                print(f"错误: 索引 {index} 对应的项目 '{info.filename}' 不是目录")
                return
            
            self.current_path = info.filename.rstrip('/')
            print(f"已进入目录: '{self.current_path}'")
                
        except Exception as e:
            print(f"更改目录时出错: {e}")

    def search_files(self, keyword):
        """根据关键字搜索文件"""
        print(f"\n在当前路径 '{self.current_path if self.current_path else '/'}' 中搜索包含 '{keyword}' 的文件:")
        print("-" * 80)
        
        try:
            infos = self.rar.infolist()
            total_files = len(infos)
            found = False
            
            for i, info in enumerate(infos):
                if keyword.lower() in info.filename.lower():
                    size = info.file_size
                    modified = self.format_time(info.date_time) if hasattr(info, 'date_time') else "N/A"
                    name = info.filename
                    
                    # 计算缩进级别以显示层次结构
                    level = name.count('/')
                    indent = "  " * level
                    
                    # 如果是目录，在名称后添加斜杠
                    if info.isdir():
                        name = name.rstrip('/') + '/'
                    
                    print(f"{i:<5} {indent}{name:<{55-2*level}} {size:<15} {modified:<20}")
                    found = True
                
                # 每100个文件显示一次进度
                if (i + 1) % 100 == 0:
                    print(f"已搜索 {i + 1}/{total_files} 个项目...")
            
            if not found:
                print("未找到匹配的文件。")
        except Exception as e:
            print(f"搜索文件时出错: {e}")
    
    def get_file_info(self, filename):
        """获取特定文件的详细信息"""
        try:
            info = self.rar.getinfo(filename)
            print(f"\n文件详细信息 '{filename}':")
            print("-" * 40)
            print(f"名称: {info.filename}")
            print(f"大小: {info.file_size} 字节")
            print(f"压缩大小: {info.compress_size} 字节")
            print(f"是否为目录: {'是' if info.isdir() else '否'}")
            if hasattr(info, 'date_time'):
                print(f"修改时间: {self.format_time(info.date_time)}")
            print(f"CRC: {info.CRC}")
            print(f"压缩方法: {info.compress_type}")
        except KeyError:
            print(f"文件 '{filename}' 不存在。")
        except Exception as e:
            print(f"获取文件信息时出错: {e}")
    
    def extract_file(self, filename, output_dir="."):
        """提取特定文件"""
        try:
            print(f"正在提取文件 '{filename}'...")
            self.rar.extract(filename, output_dir)
            print(f"文件 '{filename}' 已提取到 '{output_dir}'")
        except KeyError:
            print(f"文件 '{filename}' 不存在。")
        except Exception as e:
            print(f"提取文件时出错: {e}")

    def extract_selected_files(self, indices, output_dir="."):
        """根据序号提取选定的文件"""
        try:
            infos = self.rar.infolist()
            
            # 过滤出当前路径下的项目
            filtered_infos = []
            dir_prefix = self.current_path + "/" if self.current_path else ""
            
            for info in infos:
                # 如果在根目录，只显示顶层项目
                if not self.current_path:
                    if info.filename.count('/') == 0 or (info.filename.count('/') == 1 and info.isdir()):
                        filtered_infos.append(info)
                # 如果在子目录，只显示该目录下的直接子项目
                else:
                    if info.filename.startswith(dir_prefix):
                        relative_path = info.filename[len(dir_prefix):]
                        # 只包含直接子级（不包含子目录的子级）
                        if relative_path and (relative_path.count('/') == 0 or (relative_path.count('/') == 1 and info.isdir())):
                            filtered_infos.append(info)
            
            total_files = len(filtered_infos)
            
            # 验证序号
            valid_indices = []
            for idx in indices:
                if 0 <= idx < total_files:
                    valid_indices.append(idx)
                else:
                    print(f"警告: 序引 {idx} 超出范围，已忽略")
            
            if not valid_indices:
                print("没有有效的文件序号")
                return
            
            print(f"正在提取 {len(valid_indices)} 个文件...")
            
            # 提取选定的文件
            success_count = 0
            failed_count = 0
            
            for idx in valid_indices:
                filename = filtered_infos[idx].filename
                try:
                    print(f"正在提取文件 '{filename}'...")
                    self.rar.extract(filename, output_dir)
                    print(f"文件 '{filename}' 提取完成")
                    success_count += 1
                except Exception as e:
                    print(f"提取文件 '{filename}' 时出错: {e}")
                    failed_count += 1
            
            print(f"提取操作完成，成功 {success_count} 个文件，失败 {failed_count} 个文件")
        except Exception as e:
            print(f"批量提取文件时出错: {e}")
    
    def show_help(self):
        """显示帮助信息"""
        print("可用命令:")
        print("  list [起始索引] [数量]                      - 列出文件（默认从0开始，显示50个）")
        print("  tree [目录索引] [模式(head/all/tail/size)] [数量] - 显示指定目录下的树形结构")
        print("  cd [目录索引]                               - 进入指定索引的目录")
        print("  cd                                          - 返回上级目录")
        print("  cd -1                                       - 返回根目录")
        print("  search <关键字>                             - 搜索文件")
        print("  info <文件名>                               - 显示文件详细信息")
        print("  extract <文件名>                            - 提取文件")
        print("  extract-selected <序号列表> [-o=输出目录]    - 提取选定序号的文件（例如: extract-selected 1 5 10 -o=/tmp）")
        print("  help 或 h                                   - 显示此帮助信息")
        print("  quit 或 q                                   - 退出")

    def interactive_mode(self):
        """交互模式"""
        print("\nRAR文件查看器 - 交互模式")
        self.show_help()
        
        while True:
            try:
                command = input(f"\n[{self.current_path if self.current_path else '/'}] $ ").strip()
                if not command:
                    continue
                
                parts = command.split()
                cmd = parts[0].lower()
                
                if cmd in ['quit', 'q']:
                    print("退出程序。")
                    break
                elif cmd == 'help' or cmd == 'h':
                    self.show_help()
                elif cmd == 'list':
                    # 解析参数
                    start_index = 0
                    count = 50
                    
                    if len(parts) > 1:
                        try:
                            start_index = int(parts[1])
                        except ValueError:
                            print("错误: 起始索引必须是整数")
                            continue
                    
                    if len(parts) > 2:
                        try:
                            count = int(parts[2])
                            if count <= 0:
                                print("错误: 显示数量必须是正整数")
                                continue
                        except ValueError:
                            print("错误: 显示数量必须是整数")
                            continue
                    
                    self.list_files(start_index, count)
                elif cmd == 'tree':
                    # 解析tree命令参数
                    index = None
                    display_mode = "head"
                    count = 10
                    
                    # 解析参数
                    for i in range(1, len(parts)):
                        part = parts[i]
                        if part in ["head", "all", "tail", "size"]:
                            display_mode = part
                        elif part.isdigit():
                            try:
                                if index is None:  # 只取第一个数字作为索引
                                    index = int(part)
                                else:  # 其他数字作为count
                                    count = int(part)
                            except ValueError:
                                print(f"警告: 无法解析参数 '{part}'")
                        else:
                            try:
                                count = int(part)
                            except ValueError:
                                print(f"警告: 无法解析参数 '{part}'")
                    
                    self.list_tree(index, display_mode, count)
                elif cmd == 'cd':
                    # 解析cd命令参数
                    if len(parts) > 1:
                        if parts[1] == "-1":
                            self.cd(-1)
                        else:
                            try:
                                index = int(parts[1])
                                self.cd(index)
                            except ValueError:
                                print("错误: 索引必须是整数")
                    else:
                        self.cd()
                elif cmd == 'search' and len(parts) > 1:
                    self.search_files(parts[1])
                elif cmd == 'info' and len(parts) > 1:
                    self.get_file_info(' '.join(parts[1:]))
                elif cmd == 'extract' and len(parts) > 1:
                    self.extract_file(' '.join(parts[1:]))
                elif cmd == 'extract-selected' and len(parts) > 1:
                    try:
                        # 解析参数，支持指定输出目录
                        indices = []
                        output_dir = "."
                        
                        for part in parts[1:]:
                            if part.startswith("-o="):
                                output_dir = part[3:]
                            else:
                                indices.append(int(part))
                        
                        self.extract_selected_files(indices, output_dir)
                    except ValueError:
                        print("错误: 请提供有效的文件序号（整数）")
                else:
                    print("无效命令。输入 'list' 查看所有命令。")
            except KeyboardInterrupt:
                print("\n\n程序被用户中断。")
                break
            except Exception as e:
                print(f"发生错误: {e}")

def main():
    if len(sys.argv) < 2:
        print("用法: python rarview.py <rar文件路径>")
        sys.exit(1)
    
    rar_path = sys.argv[1]
    
    if not os.path.exists(rar_path):
        print(f"错误: 文件 '{rar_path}' 不存在")
        sys.exit(1)
    
    viewer = RARViewer(rar_path)
    viewer.interactive_mode()

if __name__ == "__main__":
    main()